%% osl_africa.m
%
% AfRICA - ArteFact Rejection using Independent Component Analysis
%
% Syntax: [fname_out scanner_artefacts_chan]=osl_africa(S)
% S needs to contain be produced inside OSL_AFRICA.m
% OPTIONAL inputs:
%   -  do_plots: set to 1 to output summary plots of artefact components.
%      set to 0 to switch off plotting.DEFAULT = 1.
%   -  do_topo: set to 1 to output summary topograghies of artefact components.
%      set to 0 to switch off topo plotting. DEFAULT = 1.
%   -  manual_approval: set to 1 to manually approve which artefact independent 
%      components get rejected. Set to 0 to for automated approval. DEFAULT
%      = 1.
%   -  do_kurt: set to positive integer "n" to view and select artefacts based on 
%      "n" low/high kurtosis. Set to -1 to view all tCIs in desending
%      order.
%      DEFAULT = 0.
%   -  do_mains: set to 1 to remove 50Hz mains artefacts.
%      set to 0 to leave mains components.DEFAULT = 1.
%   -  do_ecg: set to 1 to remove ECG artefact independent components.
%      set to 0 to leave ECG in. DEFAULT = 1. Requires ECG to be recorded
%      during acquisition.
%      do_blinks: set to 1 to remove blink artefact independent components.
%      set to 0 to leave blink in. DEFAULT = 1. Requires Eyetracker file to be
%      to be provided & eyetracker channel to be identified in SPM object.
%           - eyetracker_file: the .txt file with the eyetracker data.
%           - eyetracker_chan: the channel index of the eyetracker blink
%             channel in the SPM object. 
%      OR requires EOG channel to be specfied - eog_chan.
%      If [eyetracker_file & eyetracker_chan] or [eog_chan] are not provided 
%      then do_blinks defaults to 0.
%      If eyetracker_file & eyetracker_chan are not provided then do_blinks
%      defaults to 0.
% outputs
%   - bad_components: A list of the components identified as bad.
% HL+AB 231012

function bad_components = identify_artefactual_components_manual(S)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Set-Up
if isfield(S,'do_plots');do_plots=S.do_plots;else do_plots=1; end
if isfield(S,'manual_approval'); manual_approval=S.manual_approval; else manual_approval=1; end
if manual_approval; do_plots = 1; end
if isfield(S,'do_kurt');do_kurt=S.do_kurt;else do_kurt=0; end
if isfield(S,'do_mains');do_mains=S.do_mains;else do_mains=1; end
if isfield(S,'do_ecg'); do_ecg=S.do_ecg; else do_ecg=1; end
eyeblinks = [];
if isfield(S,'eyetracker_file') && isfield(S,'eyetracker_chan')
    S2.file=S.eyetracker_file;
    S2.eyechan=S.eyetracker_chan;
    S2.blink_win=[100,100];
    S2.D=S.D;
    eyeblinks = osl_tracker_read_file(S2);
    fs_eyetracker = round(1/diff(eyeblinks.t(1:2)));
    do_blinks=1;
elseif isfield(S,'eog_chan')
    do_blinks=1;
else
    do_blinks=0;
end
if isfield(S,'do_blinks') && do_blinks==1; do_blinks=S.do_blinks; else do_blinks=1; end   % changed by DM

D = spm_eeg_load(S.fname);

modalities = unique(D.chantype(find(strncmpi(S.modality,D.chantype,3))));  % changed by DM
try
    ecg = D(find(strcmp('ECG', D.chantype)),:);
    ecg = bandpass(ecg,[1 48],D.fsample);
catch
    warning('No ECG channel detected.');
    ecg=[];
end


%% Select only good data for classification

samples_of_interest=false(1,size(S.ica_res.tc,2));

% Remove bad trials/any trial structure
c=1;
for i=1:D.ntrials
    if D.reject == 0
    samples_of_interest(:,c:c+D.nsamples-1)=true;
    end
    c=c+D.nsamples;
end
% Remove bad segments
if D.ntrials==1;
    t = D.time;
    badsections = false(1,D.nsamples);
    Events = D.events;
    if ~isempty(Events)
        Events = Events(strcmp({Events.type},'BadEpoch'));
        for ev = 1:numel(Events)
            badsections = badsections | t >= Events(ev).time & t < (Events(ev).time+Events(ev).duration);
        end
    end
    samples_of_interest(1,badsections)=false;
end
if strcmp(S.modality,'EEG')   % changed by DM
chan_inds=setdiff(find(any([strcmp(D.chantype,'EEG')],1)),D.badchannels);
map_inds(find(any([strcmp(D.chantype,'EEG')],1))) = 1:numel(find(any([strcmp(D.chantype,'EEG')],1)));
else
chan_inds=setdiff(find(any([strcmp(D.chantype,'MEGMAG');strcmp(D.chantype,'MEGPLANAR');strcmp(D.chantype,'MEGGRAD')],1)),D.badchannels);
map_inds(find(any([strcmp(D.chantype,'MEGMAG');strcmp(D.chantype,'MEGPLANAR');strcmp(D.chantype,'MEGGRAD')],1))) = 1:numel(find(any([strcmp(D.chantype,'MEGMAG');strcmp(D.chantype,'MEGPLANAR');strcmp(D.chantype,'MEGGRAD')],1)));
end
%%

sm=S.ica_res.sm;
tc=S.ica_res.tc;

tc = tc(:,samples_of_interest);
sm=sm(map_inds(chan_inds),:);

num_ics = S.ica_res.ica_params.num_ics;
abs_ft=abs(fft(demean(tc(:,:),2),[],2));
freq_ax=0:(D.fsample/2)/(floor(length(abs_ft)/2)-1):(D.fsample/2);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Detect Mains components
fmax = zeros(1,num_ics); % added by DM
line_frequency = 50;
if do_mains
for i=1:num_ics
    spec=abs_ft(i,1:floor(length(abs_ft)/2));  % changed by DM
    [~, ind]=max(spec);
    fmax(i)=freq_ax(ind);
end
kurt=kurtosis(tc,[],2)-3;
%mains=find(line_frequency-1<fmax & fmax<line_frequency+1 & kurt'<-1);
mains=find(line_frequency-1<fmax & fmax<line_frequency+1 & kurt'<0);
%figure; plot(fmax);
msg = sprintf('\n%s%d%s\n%','Mains interference split over ', numel(mains), ' components.');
fprintf(msg);
else
    mains=[];
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Detect Cardiac components

if do_ecg   % changed by DM
if ~isempty(ecg)
ecg = bandpass(ecg,[1 48],D.fsample);
%pecg=smooth(ecg.^2);
pecg=(ecg.^2)';

r_cardiac=zeros(num_ics,size(pecg,2));
for i=1:num_ics
    sig=tc(i,:);
    %psig=smooth(sig.^2);
    psig=(sig.^2)';
    r_cardiac(i,:)=corr(psig,pecg(samples_of_interest));
end
max_cardiac=max(abs(r_cardiac),[],2);
cardiac=find(max_cardiac>0.15);
msg = sprintf('\n%s%d%s\n%s%f%s\n','Cardiac artefact split over ', numel(cardiac), ' components.');
fprintf(msg);
else
    cardiac=[];
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Detect Occular Components
if do_blinks
    
    lblink_r=zeros(1,num_ics);
    rblink_r=zeros(1,num_ics);
    lsac_r=zeros(1,num_ics);
    rsac_r=zeros(1,num_ics);
    eog_r=zeros(1,num_ics);
    
blinks_r=[];
if isfield(eyeblinks,'l_blink')
tmp=downsample(eyeblinks.l_blink,fs_eyetracker/D.fsample);
tmp = tmp(1:size(samples_of_interest,2));
for i=1:num_ics
    lblink_r(i)=corr(tc(i,:)',tmp(samples_of_interest));
end
blinks_r=[blinks_r; lblink_r];
end
if isfield(eyeblinks,'r_blink')
tmp=downsample(eyeblinks.r_blink,fs_eyetracker/D.fsample);
tmp = tmp(1:size(samples_of_interest,2));
for i=1:num_ics
    rblink_r(i)=corr(tc(i,:)',tmp(samples_of_interest));
end
blinks_r=[blinks_r; rblink_r];
end
if isfield(eyeblinks,'l_sac')
tmp=downsample(eyeblinks.l_sac,fs_eyetracker/D.fsample);
tmp = tmp(1:size(samples_of_interest,2));
for i=1:num_ics
    lsac_r(i)=corr(tc(i,:)',tmp(samples_of_interest));
end
blinks_r=[blinks_r; lsac_r];
end
if isfield(eyeblinks,'r_sac')
tmp=downsample(eyeblinks.r_sac,fs_eyetracker/D.fsample);
tmp = tmp(1:size(samples_of_interest,2));
for i=1:num_ics
    rsac_r(i,:)=corr(tc(i,:)',tmp(samples_of_interest));
end
blinks_r=[blinks_r; rsac_r];
end


S.eog_chan=find(strcmp(D.chantype,'EOG'));
if not(isempty(S.eog_chan))
eog = D(S.eog_chan,:,:); 
c=1; 
eog_rs=zeros([size(eog,1), size(eog,2)*size(eog,3)]);
for i=1:size(eog,3)
    eog_rs(:,c:c+D.nsamples-1)=eog(:,:,i);
    c=c+D.nsamples;
end
eog=eog_rs;
peog=zeros(size(eog'));
for k=1:size(eog,1)
%peog(:,k)=smooth(eog(k,:).^2);
peog(:,k)=(eog(k,:).^2)';
end
clear eog_rs;
eog_r=zeros(num_ics,size(peog,2));
for i=1:num_ics
    sig=tc(i,:);
    %psig=smooth(sig.^2);
    psig=(sig.^2)';
    eog_r(i,:)=corr(psig,peog(samples_of_interest,:));
end
blinks_r=[blinks_r eog_r];
end

%blinks_r=sum(abs(blinks_r).*(abs(blinks_r)>0.3),1)./sum((abs(blinks_r)>0.3),1); blinks_r(isnan(blinks_r))=0;
max_blinks=max(abs(eog_r),[],2);
blinks = find(max_blinks>0.15);

msg = sprintf('\n%s%d%s\n%s%f\n%s%f\n%s%f\n%s%f\n', 'Ocular artefacts split over ', numel(blinks), ' components.',...
                                                    'Right eye blink correlation = ',   abs(rblink_r(abs(rblink_r)>0.3)),...
                                                    'Left eye blink correlation = ',    abs(lblink_r(abs(lblink_r)>0.3)),...
                                                    'Right eye saccade correlation = ', abs(rsac_r(abs(rsac_r)>0.3)),...
                                                    'Left eye saccade correlation = ',  abs(lsac_r(abs(lsac_r)>0.3)));
fprintf(msg);

else
    blinks=[];
end

confirmed_artefacts=[mains ; cardiac ; blinks];
confirmed_artefacts=unique(confirmed_artefacts);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Standard plots for artefacts

if do_plots
    if ~isempty(confirmed_artefacts)
    figure; %figure('Position',[100 200*length(confirmed_artefacts) 750 200*length(confirmed_artefacts)])        
        if ~isempty(mains)
            c=1;
            for i=mains
                subplot(length(confirmed_artefacts),1+numel(modalities),c);
                plot(freq_ax,abs_ft(i,1:floor(length(abs_ft)/2)));title([{'Mains Artefacts'} {'Frequency Spectrum'}]);xlabel('Frequency (Hz)');ylabel('Spectrum');axis tight;
                c=c+1;
                for m =1:numel(modalities)
                    subplot(length(confirmed_artefacts),1+numel(modalities),c); component_topoplot(D,sm(:,i),modalities{m});
                    title([{'Mains Artefacts'} modalities(m)]);axis tight;
                    c=c+1;
                end
            end            
        end
        if ~isempty(cardiac)
            for i=cardiac
                subplot(length(confirmed_artefacts),1+numel(modalities),c);
                plot(D.time(samples_of_interest),tc(i,:));title([{'ECG Artefacts - correlation with ECG: '} {num2str(r_cardiac(i))}]);xlabel('Time (s)');axis tight;
                c=c+1;
                for m =1:numel(modalities)
                    subplot(length(confirmed_artefacts),1+numel(modalities),c); component_topoplot(D,sm(:,i),modalities{m});axis tight;
                    title([{'ECG Artefacts'} modalities(m)]);
                    c=c+1;
                end
            end
        end
        if ~isempty(blinks)
            for i=blinks
                subplot(length(confirmed_artefacts),1+numel(modalities),c);
                plot(D.time(samples_of_interest),normalise(tc(i,:)));title([{'Blink Artefacts - correlation with Eyetracker/EOG blinks: '} {num2str(blinks_r(i))}]);xlabel('Time (s)');
                if isfield(eyeblinks,'r_blink'); ho; plot(D.time,normalise(downsample(eyeblinks.r_blink,fs_eyetracker/D.fsample)),'r--');end
                if isfield(eyeblinks,'l_blink'); ho; plot(D.time,normalise(downsample(eyeblinks.l_blink,fs_eyetracker/D.fsample)),'g--');end; axis tight;
                c=c+1;
                for m =1:numel(modalities)
                    subplot(length(confirmed_artefacts),1+numel(modalities),c);component_topoplot(D,sm(:,i),modalities{m});axis tight;
                    title([{'Blink Artefacts'} modalities(m)]);
                    c=c+1;
                end
            end
        end
    end
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Manual Approval of Artefact Removal

if manual_approval
    fprintf('\n\n%s\n\n','%%%-----------------------------------------------------------%%%');
    selchan=input('Please approve artefacts for removal  (e.g. [1 3 5] to select subset, -1 to select all for removal, enter to keep all artefacts):     ');
    if ~isempty(selchan)
        if selchan==-1;
            confirmed_artefacts=confirmed_artefacts;
        else
            confirmed_artefacts=confirmed_artefacts(selchan);
        end
    else
        confirmed_artefacts=[];
    end
    close;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Artefact rejection using kurtosis.

reject_kurt=[];
if do_kurt > 0  
    comps2plot=do_kurt;

    [comps2reject] = rank_and_plot(tc,sm,abs_ft,freq_ax,D, comps2plot,'kurtosis','ascend',modalities,samples_of_interest);
    reject_kurt=[reject_kurt comps2reject];
    
    [comps2reject] = rank_and_plot(tc,sm, abs_ft,freq_ax,D, comps2plot,'kurtosis','descend',modalities,samples_of_interest);
    reject_kurt=[reject_kurt comps2reject];
    
elseif do_kurt==-1
    [comps2reject] = rank_and_plot(tc,sm,abs_ft,freq_ax,D, 'all','kurtosis','descend',modalities,samples_of_interest);
    reject_kurt=[reject_kurt comps2reject]; 
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% OUTPUT BAD COMPONENTS

bad_components=unique([confirmed_artefacts reject_kurt]);

end



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% rank_and_plot
% Subfunction to rank tICs by some summary metric and then manually flag
% artefact components from topoplots, time course and spectrum.

function [comps2reject] = rank_and_plot(tc,sm,abs_ft,freq_ax,D,comps2plot,metric,direction,modalities,samples_of_interest)

switch lower(metric)
    case 'kurtosis'
        met=kurtosis(tc,[],2)-3;
        figlab='Kurtosis';
    otherwise
end

[met_ord, order]=sort(met,direction);
comps2reject=[];

if ~strcmp(comps2plot,'all')
comp=sm(:,order(1:comps2plot));
fh(1)=figure;%('Position',[100 200*comps2plot 750 200*comps2plot]);
for i = 1:size(comp,2)   
    for m = 1:numel(modalities)
        subplot(size(comp,2),2+numel(modalities),(2+numel(modalities))*(i-1)+m);component_topoplot(D,comp(:,i),modalities{m}) ;title({['Component ' num2str(order(i)) ': ' figlab ' = ' num2str(met_ord(i))] modalities{m}}); axis tight;
    end
    subplot(size(comp,2),2+numel(modalities),(2+numel(modalities))*(i-1)+m+1); plot(D.time(samples_of_interest),tc(order(i),:)); title({['Component ' num2str(order(i)) ': ' figlab ' = ' num2str(met_ord(i))] 'Independent Time Course'}); xlabel('Time (s)');axis tight;
    subplot(size(comp,2),2+numel(modalities),(2+numel(modalities))*(i-1)+m+2); plot(freq_ax,abs_ft(order(i),1:floor(length(abs_ft)/2))); title({['Component ' num2str(order(i)) ': ' figlab ' = ' num2str(met_ord(i))] 'Frequency Spectrum'}); xlabel('Frequency (Hz)');axis tight;
end

fprintf('\n\n%s\n\n','%%%-----------------------------------------------------------%%%');
selchan=input(['Please approve artefact components based on ' metric ' for removal  (e.g. [1 3 5] to select subset, -1 to select all for removal, enter to keep all artefacts):     ']);
if ~isempty(selchan)
    if selchan==-1;
        comps2reject=order(1:comps2plot)';
    else
        comps2reject=order(selchan)';
    end
else
    comps2reject=[];
end

else
    fh=figure;%('Position',[100 600 600 600]) ;
    comp=sm(:,order);
    
    i=1;exit_flag=0;
    while i <= size(comp,2) && exit_flag==0
        subplot(3,1,1); plot(D.time(samples_of_interest),tc(order(i),:)); title({['Component ' num2str(order(i)) ': ' figlab ' = ' num2str(met_ord(i))] 'Independent Time Course'}); xlabel('Time (s)');
        subplot(3,1,2); plot(freq_ax,abs_ft(order(i),1:floor(length(abs_ft)/2)));title({['Component ' num2str(order(i)) ': ' figlab ' = ' num2str(met_ord(i))] 'Frequency Spectrum'}); xlabel('Frequency (Hz)');
        for m = 1:numel(modalities)
        subplot(3,numel(modalities),2*numel(modalities)+m);component_topoplot(D,comp(:,i),modalities{m}) ;title({['Component ' num2str(order(i)) ': ' figlab ' = ' num2str(met_ord(i))] modalities{m}}); axis tight;
        end
        fprintf('\n\n%s\n\n','%%%-----------------------------------------------------------%%%');
        selchan=input(['Please approve artefact components based on ' metric ' for removal  (e.g. 1 to reject, enter to leave in, -1 to skip remaining) :     ']);
        if selchan == 1
            comps2reject=[comps2reject order(i)];
        elseif selchan == -1
            exit_flag=1;
        end
        clf(fh);
        i=i+1;
    end
end
close(fh);
end

%% function to produce component topoplot

function component_topoplot(D,comp,modality)
cfg=[];
data=[];

global OSLDIR

chan_inds=setdiff(find(any([strcmp(D.chantype,'EEG');strcmp(D.chantype,'MEGMAG');strcmp(D.chantype,'MEGPLANAR');strcmp(D.chantype,'MEGGRAD')],1)),D.badchannels);
map_inds(find(any([strcmp(D.chantype,'EEG');strcmp(D.chantype,'MEGMAG');strcmp(D.chantype,'MEGPLANAR');strcmp(D.chantype,'MEGGRAD')],1))) = 1:numel(find(any([strcmp(D.chantype,'EEG');strcmp(D.chantype,'MEGMAG');strcmp(D.chantype,'MEGPLANAR');strcmp(D.chantype,'MEGGRAD')],1)));


comp_full(map_inds(chan_inds),:)=comp;
comp_full(D.badchannels) = 0;
comp = comp_full;

if (strcmp(modality,'MEGMAG')),
    cfg.channel     = {'MEGMAG'};
    cfg.layout      = [OSLDIR '/layouts/neuromag306mag.lay'];
    chanind = strmatch(cfg.channel, D.chantype);
    
    comp2view=zeros(102,1);
    for ind=1:length(chanind),
        indp=(ind-1)*3+3;
        comp2view(ind,:)=comp(indp);
    end;
    data.dimord='chan_comp';
    data.topo=comp2view;
    data.topolabel = D.chanlabels(chanind);
    data.time = {1};
    
elseif (strcmp(modality,'MEGPLANAR')),
    cfg.channel     = {'MEGMAG'};
    cfg.layout      = [OSLDIR '/layouts/neuromag306mag.lay'];
    chanind = strmatch(cfg.channel, D.chantype);
    comp2view=zeros(102,1);
    for ind=1:length(chanind),
        indp=(ind-1)*3+1;
        comp2view(ind,:)=comp(indp)+comp(indp+1);
    end;
    data.dimord='chan_comp';
    data.topo=comp2view;
    data.topolabel = D.chanlabels(chanind);
    data.time = {1};
    warning('The planar grads are being overlayed on the magnetometer locations and using the magnetometer labels');

elseif (strcmp(modality,'MEGGRAD')),
    cfg.channel     = {'MEG'};
    cfg.layout      = [OSLDIR '/layouts/CTF275.lay'];
    chanind = strmatch(cfg.channel, D.chantype);
    comp2view=comp;
    data.dimord='chan_comp';
    data.topo=comp2view;
    data.topolabel = D.chanlabels(chanind);
    data.time = {1};
    
elseif (strcmp(modality,'EEG')),
    cfg.channel     = {'EEG'};
    A=sensors(D,'EEG');
    pnt=A.chanpos;
    elec = [];
    elec.pnt = pnt(chan_inds,:);
    elec.label = A.label(chan_inds);
    cfg.elec  = elec;
    %cfg.layout = ft_prepare_layout(cfg);
    cfg.channel={'all'};
    
    comp2view=comp;
    data.dimord='chan_comp';
    data.topo=comp2view;
    data.topolabel = A.label(chan_inds);
    data.label = A.label(chan_inds);
    data.time = {1};
else
    
    error('Unsupported modality');
end

cfg.component = 1;
cfg.interactive = 'no';
cfg.comment     = '';
ft_topoplotIC(cfg,data);

end

